#pragma GCC optimize(2)
#pragma GCC optimize(3,"Ofast","inline")

#include "depth_to_space_tiling.h"
#include "register/op_def_registry.h"

// using namespace std;

#define ull uint64_t

namespace optiling {
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{

    DepthToSpaceTilingData tiling;

    int attr_num = context->GetAttrs()->GetAttrNum();
    // cout << "Attr Num: " << attr_num << endl;
    const int64_t *block_size0 = context->GetAttrs()->GetInt(0);
    int8_t mode = 0;  // DCR
    int8_t data_format = 0;  // NHWC
    if (attr_num > 1) {
        const char *mode0 = context->GetAttrs()->GetStr(1);
        // cout << "Mode: " << mode0 << endl;
        if (mode0[0] == 'C') {
            mode = 1;  // CRD
        }
        const char *data_format0 = context->GetAttrs()->GetStr(2);
        // cout << "Data Format: " << data_format0 << endl;
        if (data_format0[1] == 'C') {
            data_format = 1;  // NCWH
        }
    }
    ull block_size = *block_size0;
    // cout << "Block Size: " << block_size << endl;
    ull n, h, w, c, cb, cbb;
    ull dim[4];
    ull totalLength = 1;
    auto shape_x = context->GetInputShape(0)->GetOriginShape();
    // cout << "Shape Dimensions: ";
    for (int i = 0; i < 4; i++) {
        ull nowdim = shape_x.GetDim(i);
        // cout << nowdim << " ";
        totalLength *= nowdim;
        dim[i] = nowdim;
    }
    // cout << endl;
    if (data_format == 0) {
        // NHWC
        n = dim[0];
        h = dim[1];
        w = dim[2];
        c = dim[3];
    } else {
        // NCWH
        n = dim[0];
        c = dim[1];
        h = dim[2];
        w = dim[3];
    }
    cb = c / block_size;
    cbb = cb / block_size;
    // printf("mode: %d, data_format: %d\n", mode, data_format);
    // printf("block_size: %d, n: %d, h: %d, w: %d, c: %d, cb: %d, cbb: %d\n", block_size, n, h, w, c, cb, cbb);

    ull input_size = 0;
    auto tensor = context->GetInputTensor(0);
    auto type = tensor->GetDataType();
    if(type==ge::DT_FLOAT16) {
        input_size = sizeof(float) / 2;
    } else if(type==ge::DT_FLOAT) {
        input_size = sizeof(float);
    } else if(type==ge::DT_INT32) {
        input_size = sizeof(int32_t);
    } else if(type==ge::DT_INT8) {
        input_size = sizeof(int8_t);
    }

    // cout << "Input Size: " << input_size << endl;
    
    if (!(mode == 0 && data_format == 0)) {
        // case 2~4
        context->SetTilingKey(1);
    } else if (block_size == 2 && n == 3) {
        // case 1
        context->SetTilingKey(1);
    } else {
        // 性能测试
        context->SetTilingKey(2);
    }

    if (context->GetTilingKey() == 1) {
        context->SetBlockDim(1);
        tiling.set_tileNum(1);
    } else {
        context->SetBlockDim(n * h);
        tiling.set_tileNum(1);
    }
    // cout << "Block Dim: " << context->GetBlockDim() << endl;

    tiling.set_totalLength(totalLength);
    tiling.set_mode(mode);
    tiling.set_data_format(data_format);
    tiling.set_block_size(block_size);
    tiling.set_n(n);
    tiling.set_h(h);
    tiling.set_w(w);
    tiling.set_c(c);
    tiling.set_cb(cb);
    tiling.set_cbb(cbb);
    tiling.set_input_size(input_size);
    
    tiling.set_blockLength(totalLength / context->GetBlockDim());
    tiling.set_tileLength(cb);
    tiling.set_tileSize(cb * input_size);
    // DCR, NHWC
    // ull shape[6] = {n, h, w, block_size, block_size, cbb};
    // ull trans[6] = {0, 1, 3, 2, 4, 5};
    tiling.set_shape2(w);
    tiling.set_shape3(block_size);
    tiling.set_suf3(c);
    tiling.set_suf4(cb);
    tiling.set_transsuf3(cb*w);
    tiling.set_transsuf4(cb);

    tiling.set_bufferSize(cb * input_size * block_size);
    tiling.set_stride(cb * w * input_size - cb * input_size);

    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());

    return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
}


namespace ops {
class DepthToSpace : public OpDef {
public:
    explicit DepthToSpace(const char* name) : OpDef(name)
    {
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT32, ge::DT_INT8})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT32, ge::DT_INT8})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Attr("block_size").Int();
        this->Attr("mode").AttrType(OPTIONAL).String("DCR");
        this->Attr("data_format").AttrType(OPTIONAL).String("NHWC");

        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend910b");

    }
};

OP_ADD(DepthToSpace);
}
